#encoding=utf-8
'''
Created on 2011-11-14

@author: guanyw

python version:2.6
'''

from xml.etree import ElementTree

def norm(value_string):
    """
    Normalize the string value in an RTE pair's C{value} or C{entailment} 
    attribute as an integer (1, 0).
    
    @param value_string: the label used to classify a text/hypothesis pair
    @type value_string: C{str}
    @rtype: C{int}
    """

    valdict = {"Y":1,
               "YB":1,
               "N":0,
               "TRUE": 1,
               "FALSE": 0,
               "YES": 1,
               "NO": 0}
    return valdict[value_string.upper()]

    
class RTEPairSimple:
    '''RTE Pair: Simplify RTEPair'''
    
    def __init__(self, pair,doc_id=None,topic_id=None,text_id=None,
                 hyp_id=None,text=None,hyp=None,value=None):
        """
        @param doc_id: the text document id
        @param topic_id : topic id 
        @param text: the text component of the pair
        @param hyp: the hypothesis component of the pair
        @param value: classification label for the pair
        """
        def set_field(field_name,field_value):
            if field_name in pair.attrib:
                return pair.attrib[field_name]
            return field_value
        
        self.topic_id = set_field('t_id', topic_id)
        self.doc_id = set_field('doc_id', doc_id)
        self.text_id = set_field('s_id', text_id)
        self.hyp_id = set_field('h_id', hyp_id)
        self.value = set_field('value', value)
        self.value = set_field('tag', self.value)
        self.value = set_field('label', self.value)
#        self.text = set_field('t1', text)
#        self.hyp = set_field('t2', hyp)
        try:
            self.text = pair[0].text
            self.hyp = pair[1].text
        except:
            pass
        
        if self.value:
            self.value = norm(self.value)
        
        if self.text_id and not text:
            self.text = pair.text
        elif text:
            self.text = text
            
        if self.hyp_id and not hyp:
            self.hyp = pair.text.strip()
        elif hyp:
            self.hyp = hyp
    
    
    def __repr__(self):
        return '\tt_id:%s; doc_id:%s; text_id:%s; hyp_id:%s;value:%s;text:%s;hyp:%s\n' \
    %(self.topic_id,self.doc_id,self.text_id,self.hyp_id,self.value,self.text,self.hyp)
        
class RTEReaderFromTrainPairs:
    """
    Corpus reader for corpora in RTE train pairs.
    
    FORMAT:
    <pair ....>
        <t1>...</t1>
        <t2>...</t2>
    </pair>
    """
    def xml(self, fileid):
        if not isinstance(fileid, basestring):
            raise TypeError('Expected a single file identifier string')
        # Read the XML in using ElementTree.
        return ElementTree.parse(fileid)
    
    def _read_etree(self, doc):
        """
        Map the XML input into an RTEPair.
        
        This uses the C{getiterator()} method from the ElementTree package to
        find all the C{<pair>} elements.
        
        @param doc: a parsed XML document
        @rtype: C{list} of L{RTEPair}s
        """
        return [RTEPairSimple(pair) for pair in doc.getiterator("pair")]

    def pairs(self, fileids):
        """
        Build a list of RTEPairs from a RTE corpus.
        
        @param fileids: a list of RTE corpus fileids
        @type: C{list}
        @rtype: C{list} of L{RTEPair}s
        """
        if isinstance(fileids, basestring): fileids = [fileids]
        pair_list = []
        for fileid in fileids:
            pair_list.extend(self._read_etree(self.xml(fileid))) 
        return  pair_list
    


class RTERederFromHypothesis(RTEReaderFromTrainPairs):
    """
    Corpus reader for RTE corpus.
    
    
    TWO FORMAT: 
    hypothesis :
    <TOPIC t_id="1">
        <H h_id="11">...</H>
        <H h_id="12">...</H>
    </TOPIC>
    <TOPIC t_id="2">
        <H h_id="21">...</H>
    </TOPIC>
    
    text :
    <document>
        <s s_id="0">...</s>
        <s s_id="0">...</s>
    </document>
    """
    def __init__(self,root):
        self.root = root
        
    def _read_etree(self, doc):
        """
        Map the XML input into an RTEPair.
        
        This uses the C{getiterator()} method from the ElementTree package to
        find all the C{<pair>} elements.
        
        @param doc: a parsed XML document
        @rtype: C{list} of L{RTEPair}s
        """
        for pair in doc.getiterator('s'):
            if 's_id' in pair.attrib:

                return [RTEPairSimple(pair) for pair in doc.getiterator("s")]
            
        ret_list = []
        for pair in doc.getiterator('TOPIC'):
            topic_id = pair.attrib['t_id']
            for hyp_pair in pair:

                for doc_id in range(1,11):
                    text_pair = self.pairs(self.root+'topics/%s/%s.xml'%(topic_id,doc_id))
                                        
                    for tp in text_pair:
                        ret_list.append(RTEPairSimple(hyp_pair,
                                                      topic_id=topic_id,
                                                      doc_id=doc_id,
                                                      text=tp.text,
                                                      text_id=tp.text_id))
            
        return ret_list


class RTEReaderFrom1ToNPairs(RTEReaderFromTrainPairs):
    """
    Corpus reader for RTE corpus.
    
    
    FORMAT: 
    <entailment_corpus>
      <TOPIC t_id="1">
        <H h_id="11">
          <H_sentence>hypothesis</H_sentence>
          <text s_id="0" doc_id="1" value="N">text</text>
          <text s_id="2" doc_id="1" value="N">text</text>
            ...
        </H>
      </TOPIC>
      <TOPIC t_id="2">
        <H h_id="21">
          <H_sentence>hypothesis</H_sentence>
          <text s_id="0" doc_id="1" value="N">text</text>
            ...
        </H>
      </TOPIC>
    </entailment_corpus>
    """
        
    def _read_etree(self, doc):
        """
        Map the XML input into an RTEPair.
        
        This uses the C{getiterator()} method from the ElementTree package to
        find all the C{<pair>} elements.
        
        @param doc: a parsed XML document
        @rtype: C{list} of L{RTEPairSimple}s
        """
            
        ret_list = []
        for pair in doc.getiterator('TOPIC'):
            topic_id = pair.attrib['t_id']
            for hyp_pair in pair:
                hyp_id = hyp_pair.attrib['h_id']
                hyp = hyp_pair[0].text

                if len(hyp_pair) >1:
                    for pair in hyp_pair[1:]:
                        ret_list.append(RTEPairSimple(pair,topic_id=topic_id,hyp_id=hyp_id,hyp=hyp))
        
        return ret_list

if __name__ == '__main__':
    pass
    
